package cn.itcast.tags.models.ml

import cn.itcast.tags.models.{AbstractModel, ModelType}
import cn.itcast.tags.tools.TagTools
import org.apache.spark.ml.clustering.{KMeans, KMeansModel}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.{Column, DataFrame}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{DataTypes, DecimalType}

import scala.tools.nsc.interpreter.session

/**
 * 挖掘类型标签模型开发：用户客户价值标签模型（RFM模型）
 */
class RfmModel extends AbstractModel("客户价值标签", ModelType.ML) {
  /*
    361	客户价值
      362	高价值		  0
      363	中上价值		1
      364	中价值		  2
      365	中下价值		3
      366	超低价值		4
   */
  override def doTag(businessDF: DataFrame, tagDF: DataFrame): DataFrame = {
    val session = businessDF.sparkSession
    import session.implicits._

    // businessDF.printSchema()
    // businessDF.show(10, truncate = false)

    // tagDF.printSchema()
    // tagDF.filter($"level" === 5).show(10, truncate = false)

    /*
      TODO: 1、从订单数据获取字段值，计算每个用户RFM值
        按照memberId分组，然后进行聚合计算
        R: 最后一次消费时间距离今天的天数
          日期函数：current_timestamp(), from_unixtime(), datediff()
        F: 消费次数（订单次数）
          count函数
        M：消费金额（所有订单总的消费金额）
          sum函数
     */
    val rfmDF: DataFrame = businessDF
      // 按照用户进行分组
      .groupBy($"memberid")
      // 聚合计算R、F、M值
      .agg(
        max($"finishtime").as("max_finish_time"),
        count($"ordersn").as("frequency"),
        sum($"orderamount".cast(DataTypes.createDecimalType(10, 2))).as("monetary")
      )
      .select(
        $"memberid".as("userId"),
        datediff(current_timestamp(), from_unixtime($"max_finish_time")).as("recency"),
        $"frequency", $"monetary"
      )
    /*
      root
       |-- userId: string (nullable = true)
       |-- recency: integer (nullable = true)
       |-- frequency: long (nullable = false)
       |-- monetary: decimal(20,2) (nullable = true)
     */
    // rfmDF.printSchema()
    // rfmDF.show(10, truncate = false)

    /*
      TODO: 2、按照规则，给RFM值打分Score
        R: 1-3天=5分，4-6天=4分，7-9天=3分，10-15天=2分，大于16天=1分
        F: ≥200=5分，150-199=4分，100-149=3分，50-99=2分，1-49=1分
        M: ≥20w=5分，10-19w=4分，5-9w=3分，1-4w=2分，<1w=1分
     */
    // R 打分判断条件
    val rWhen: Column = when($"recency".between(1, 3), 5.0)
      .when($"recency".between(4, 6), 4.0)
      .when($"recency".between(7, 9), 3.0)
      .when($"recency".between(10, 15), 2.0)
      .when($"recency".gt(16), 1.0)
    // F 打分条件表达式
    val fWhen = when(col("frequency").between(1, 49), 1.0)
      .when(col("frequency").between(50, 99), 2.0)
      .when(col("frequency").between(100, 149), 3.0)
      .when(col("frequency").between(150, 199), 4.0)
      .when(col("frequency").geq(200), 5.0)
    // M 打分条件表达式
    val mWhen = when(col("monetary").lt(10000), 1.0)
      .when(col("monetary").between(10000, 49999), 2.0)
      .when(col("monetary").between(50000, 99999), 3.0)
      .when(col("monetary").between(100000, 199999), 4.0)
      .when(col("monetary").geq(200000), 5.0)
    // 基于规则对RFM值进行打分
    val rfmScoreDF: DataFrame = rfmDF
      .select(
        $"userId",
        rWhen.as("r_score"),
        fWhen.as("f_score"),
        mWhen.as("m_score")
      )
    // rfmScoreDF.printSchema()
    // rfmScoreDF.show(50, truncate = false)

    // TODO: 3、将RFM数据进行聚类（KMeans算法、K=5）
    // 3.1 组合R、F、M到特征向量vector中
    val assembler: VectorAssembler = new VectorAssembler()
      .setInputCols(Array("r_score", "f_score", "m_score"))
      .setOutputCol("features")
    val featuresDF: DataFrame = assembler.transform(rfmScoreDF)
    // 将数据进行缓存
    featuresDF.cache().count()

    // 3.2 使用KMeans算法聚类，训练模型
    /*
      val kMeansModel: KMeansModel = new KMeans()
        .setFeaturesCol("features")
        .setPredictionCol("prediction")
        // K值得设置
        .setK(5)
        .setMaxIter(20)
        // 训练模型
        .fit(featuresDF)
      // println(s"WSSSE = ${kMeansModel.computeCost(featuresDF)}")
     */
    val kMeansModel: KMeansModel = trainModel(featuresDF)

    // 3.3 使用模型对数据进行预测：划分类簇，属于哪个类别
    val predictionDF: DataFrame = kMeansModel.transform(featuresDF)
    /*
      root
       |-- userId: string (nullable = true)
       |-- r_score: double (nullable = true)
       |-- f_score: double (nullable = true)
       |-- m_score: double (nullable = true)
       |-- features: vector (nullable = true)
       |-- prediction: integer (nullable = true)
     */
    // predictionDF.printSchema()
    // predictionDF.show(50, truncate = false)

    // 3.4 获取类簇中心点，采用拉链操作，获取索引
    val centerIndexArray: Array[((Int, Double), Int)] = kMeansModel
      .clusterCenters
      .zipWithIndex
      // 获取类簇中心点向量之和（RFM之和）
      .map { case (vector, clusterIndex) => (clusterIndex, vector.toArray.sum) }
      // 对RFM降序排序
      .sortBy { case (_, rfm) => -rfm }
      // 再次拉链获取索引，匹配定义好的规则
      .zipWithIndex
    // centerIndexArray.foreach(println)

    // TODO: 4、从KMeans中获取出每个用户属于簇
    // 4.1 获取属性标签（5级标签）数据，选择rule和tagName
    val rulesMap: Map[String, String] = TagTools.convertMap(tagDF)
    // rulesMap.foreach(println)
    // println("===================================")

    // 4.2 聚类类簇关联属性标签数据rule，对应聚类类簇与标签tagName
    val indexTagMap: Map[Int, String] = centerIndexArray
      .map { case ((centerIndex, _), index) =>
        val tagName = rulesMap(index.toString)
        (centerIndex, tagName)
      }
      .toMap
    indexTagMap.foreach(println)

    // 4.3 自定义UDF函数，传递prediction，返回tagName
    val indexTagMapBroadcast = session.sparkContext.broadcast(indexTagMap)
    val index_to_tag = udf(
      (prediction: Int) => indexTagMapBroadcast.value(prediction)
    )

    // 4.4 对预测数据，打标签
    val modelDF: DataFrame = predictionDF
      .select(
        $"userId",
        index_to_tag($"prediction").as("rfm")
      )
    // modelDF.show(100, truncate = false)

    // 5. 返回标签模型数据
    modelDF
  }

  /**
   * 使用KMeans算法训练模型
   *
   * @param dataframe 数据集
   * @return KMeansModel模型
   */
  def trainModel(dataframe: DataFrame): KMeansModel = {
    // 构建KMeans算法对象，设置相关参数
    val kMeansModel: KMeansModel = new KMeans()
      .setFeaturesCol("features")
      .setPredictionCol("prediction")
      .setK(5) // 设置列簇个数：5
      .setMaxIter(10) // 设置最大迭代次数
      .fit(dataframe)
    println(s"WSSSE = ${kMeansModel.computeCost(dataframe)}")
    // 返回模型对象
    kMeansModel
  }
}

object RfmModel {
  def main(args: Array[String]): Unit = {
    val tagModel = new RfmModel()
    tagModel.executeModel(361L)
  }
}